package com.ccx.myinterpreter;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Stack;

import com.ccx.myinterpreter.Expr.Assignment;
import com.ccx.myinterpreter.Expr.Binary;
import com.ccx.myinterpreter.Expr.Call;
import com.ccx.myinterpreter.Expr.Get;
import com.ccx.myinterpreter.Expr.Grouping;
import com.ccx.myinterpreter.Expr.Literal;
import com.ccx.myinterpreter.Expr.Set;
import com.ccx.myinterpreter.Expr.Super;
import com.ccx.myinterpreter.Expr.Unary;
import com.ccx.myinterpreter.Expr.Variable;
import com.ccx.myinterpreter.Stmt.Block;
import com.ccx.myinterpreter.Stmt.ClassDecl;
import com.ccx.myinterpreter.Stmt.Expression;
import com.ccx.myinterpreter.Stmt.FunDecl;
import com.ccx.myinterpreter.Stmt.IfStmt;
import com.ccx.myinterpreter.Stmt.Print;
import com.ccx.myinterpreter.Stmt.ReturnStmt;
import com.ccx.myinterpreter.Stmt.VarDecl;
import com.ccx.myinterpreter.Stmt.WhileStmt;
import com.ccx.myinterpreter.Stmt.forStmt;

public class Resolver implements Expr.Visitor, Stmt.Visitor {

    private Interpreter interpreter;
    private Stack<Map<String, Boolean>> stack;

    public Resolver(Interpreter interpreter) {
        this.interpreter = interpreter;
        stack = new Stack<>();
    }

    public void resolve(List<Stmt> stmts) {
        for (Stmt stmt : stmts) {
            stmt.accept(this);
        }
    }

    private void resolveVariable(Expr variable, String name) {
        for (int i = stack.size() - 1; i >= 0; i--) {
            if(stack.get(i).containsKey(name)) {
                interpreter.varBind(variable, stack.size() - 1 -i);
            }
        }
    }

    private void resolve(Stmt stmt) {
        stmt.accept(this);
    }

    private void resolveFunDecl(FunDecl funDecl) {
        startScope();
        for (Token param : funDecl.params) {
            declear(param.getLexeme());
            define(param.getLexeme());
        }
        for (Stmt stmt : funDecl.body.stmts) {
            resolve(stmt);
        }
        endScope();
    }

    private void startScope() {
        stack.push(new HashMap<>());
    }

    private void endScope() {
        stack.pop();
    }

    private void define(String key) {
        if(stack.isEmpty()) return;
        stack.peek().put(key, true);
    }

    private void resolve(Expr expr) {
        expr.accept(this);
    }

    private void declear(String key) {
        if(stack.isEmpty()) return;
        if(stack.peek().containsKey(key))
            error("Variable " + key + " already exist!");
        stack.peek().put(key, false);
    }

    private void error(String msg) {
        throw new RuntimeException(msg);
    }

    @Override
    public void visitExpr(Expression stmt) {
        stmt.expr.accept(this);
    }

    @Override
    public void visitPrint(Print stmt) {
        stmt.expr.accept(this);
    }

    @Override
    public void visitVarDecl(VarDecl stmt) {
        declear(stmt.identifier.getLexeme());
        if(stmt.expr != null) {
            resolve(stmt.expr);
        }
        define(stmt.identifier.getLexeme());
    }

    @Override
    public void visitBlock(Block block) {
        startScope();
        for (Stmt stmt : block.stmts) {
            resolve(stmt);
        }
        endScope();
    }

    @Override
    public void visitIf(IfStmt ifStmt) {
        ifStmt.expr.accept(this);
        ifStmt.thenStmt.accept(this);
        if(ifStmt.elseStmt != null) ifStmt.elseStmt.accept(this);
    }

    @Override
    public void visitWhileStmt(WhileStmt whileStmt) {
        whileStmt.condiation.accept(this);
        whileStmt.stmt.accept(this);
    }

    @Override
    public void visitForStmt(forStmt forStmt) {
        if(forStmt.first != null) forStmt.first.accept(this);
        if(forStmt.second != null) forStmt.second.accept(this);
        if(forStmt.third != null) forStmt.third.accept(this);
        forStmt.body.accept(this);
    }

    @Override
    public void visitFunDecl(FunDecl funDecl) {
        declear(funDecl.name.getLexeme());
        define(funDecl.name.getLexeme());
        resolveFunDecl(funDecl);
    }

    @Override
    public void visitReturnStmt(ReturnStmt returnStmt) {
        if(returnStmt.expr != null) {
            returnStmt.expr.accept(this);
        }
    }

    @Override
    public void visitClassDecl(ClassDecl classDecl) {
        declear(classDecl.name.getLexeme());
        define(classDecl.name.getLexeme());
        if(classDecl.superClass != null) {
            classDecl.superClass.accept(this);
            if(classDecl.name.getLexeme().equals(classDecl.superClass.name.getLexeme())) {
                throw new RuntimeException("class can't extends itself!");
            }
            startScope();
            stack.peek().put("super", true);
        }
        startScope();
        stack.peek().put("this", true);
        for (FunDecl method : classDecl.methods) {
            resolveFunDecl(method);
        }
        endScope();
        if(classDecl.superClass != null) {
            endScope();
        }
    }

    @Override
    public Object visitLiteral(Literal literal) {
        return null;
    }

    @Override
    public Object visitGrouping(Grouping grouping) {
        grouping.expr.accept(this);
        return null;
    }

    @Override
    public Object visitUnary(Unary unary) {
        unary.right.accept(this);
        return null;
    }

    @Override
    public Object visitBinary(Binary binary) {
        binary.left.accept(this);
        binary.right.accept(this);
        return null;
    }

    @Override
    public Object visitVariable(Variable variable) {
        if(!stack.isEmpty() && stack.peek().get(variable.name.getLexeme()) == Boolean.FALSE) {
            error("Can't read variable in it's initelizer");
        }
        resolveVariable(variable, variable.name.getLexeme());
        return null;
    }

    @Override
    public Object visitAssignment(Assignment assignment) {
        resolveVariable(assignment, assignment.name.getLexeme());
        return null;
    }

    @Override
    public Object visitCall(Call call) {
        call.callee.accept(this);
        for (Expr arg : call.args) {
            arg.accept(this);
        }
        return null;
    }

    @Override
    public Object visitGet(Get get) {
        get.object.accept(this);
        return null;
    }

    @Override
    public Object visitSet(Set set) {
        set.object.accept(this);
        set.value.accept(this);
        return null;
    }

    @Override
    public Object visitSuper(Super superExpr) {
        resolveVariable(superExpr, "super");
        return null;
    }
    
}
